from .dukemtmc import DukeMTMC
from .market1501 import Market1501
from .msmt import MSMT


__all__ = ['DukeMTMC','Market1501','MSMT']



def collate_fn(batch):
    images = []
    person_ids = []
    camera_ids = []
    video_ids = []
    filenames = []
    for img,pid,camid,vid,filename in batch:
        images.append(img)
        person_ids.append(pid)
        camera_ids.append(camid)
        video_ids.append(vid)
        filenames.append(filename)
    return torch.stack(images,dim=0),torch.tensor(person_ids).long(),torch.tensor(camera_ids),torch.tensor(video_ids),filenames





